import matplotlib.pyplot as plt #matplotlib绘图
from mpl_toolkits.mplot3d import Axes3D #plt3D绘图
import torch   #深度学习，conda环境下pytorch或在命令行安装
import copy  #copy模块拷贝对象

#Lorenz吸引子生成函数,参数为三个初始坐标，三个初始参数,迭代次数,返回三个一维list
def lorenz_attractor(x_0, y_0, z_0, sigma=10, beta=8/3, rho=28, steps=10000, end=10):
    x = [x_0]
    y = [y_0]   #返回的三维列表  x,y,z
    z = [z_0]

    for _ in range(steps):   #对10000次步长遍历  根据公式计算
        x.append(x[-1] + (sigma * (y[-1] - x[-1])) * end / steps)
        y.append(y[-1] + (x[-1] * (rho - z[-1]) - y[-1]) * end / steps)
        #rho 为ρ，二维正态分布参数，在定义吸引子时设置
        z.append(z[-1] + (x[-1] * y[-1] - beta * z[-1]) * end / steps)
    return (x, y, z)  #依次迭代计算返回坐标值


def lorenz_attractor_endpoint(x_0, y_0, z_0, sigma=10, beta=8 / 3, rho=28, steps=10000, end=10):
    x = copy.deepcopy(x_0)    #同样的方法copy方法记录坐标变换
    y = copy.deepcopy(y_0)
    z = copy.deepcopy(z_0)
    for _ in range(steps):
        x = x + sigma * (y - x) * end / steps
        y = y + (x * (rho - z) - y) * end / steps
        z = z + (x * y - beta * z) * end / steps
    return (x, y, z)

if __name__ == "__main__":     #使用matplotlib的3D模块绘图
    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")

    x0 = 5.0   #给x,y,z赋值
    y0 = 5.0
    z0 = 5.0

    x, y, z = lorenz_attractor(x0, y0, z0)
    ax.plot(x, y, z)

    x_end, y_end, z_end = lorenz_attractor_endpoint(torch.tensor([x0, x0 + 1]), torch.tensor([y0, y0]),
                                                    torch.tensor([z0, z0]))
    ax.scatter([x0, x0 + 1], [y0, y0], [z0, z0], color="green")  #起点
    ax.scatter(x_end, y_end, z_end, color="red")  #终点
    plt.show()